{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.utils.data as Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#自制数据集\n",
    "             # Encoder_input    Decoder_input        Decoder_output\n",
    "sentences = [['我 是 学 生 P' , 'S I am a student'   , 'I am a student E'],         # S: 开始符号\n",
    "             ['我 喜 欢 学 习', 'S I like learning P', 'I like learning P E'],      # E: 结束符号\n",
    "             ['我 在 北 京 P',  'S I in bei jing '   , 'I in bei jing E'],          \n",
    "             ['中 国 首 都 P',  'S bei jing P P '    , 'bei jing P P E'],          \n",
    "             ['I am a boy P',  'S 我 是 男 生 '      , '我 是 男 生 E'] ,            \n",
    "             ['我 是 男 生 P',  'S I am a boy'       , 'I am a boy E']]             # P: 占位符号，如果当前句子不足固定长度用P占位 pad补0\n",
    "\n",
    "\n",
    "src_vocab = {'P':0, '我':1, '是':2, '学':3, '生':4, '喜':5, '欢':6,'习':7,'男':8,'在':9,'北':10,'京':11,\n",
    "             '中':12,'国':13,'首':14,'都':15,'I':16,'am':17, 'a':18,'boy':19, 'E':20\n",
    "             }   # 词源字典  字：索引\n",
    "src_idx2word = {src_vocab[key]: key for key in src_vocab}\n",
    "src_vocab_size = len(src_vocab)                 # 字典字的个数\n",
    "\n",
    "tgt_vocab = {'S':0, 'E':1, 'P':2, 'I':3, 'am':4, 'a':5, 'student':6, 'like':7, 'learning':8, 'boy':9,\n",
    "             'in':10, 'bei':11, 'jing':12, '我':13, '是':14,'男':15,'生':16\n",
    "             }\n",
    "idx2word = {tgt_vocab[key]: key for key in tgt_vocab}                               # 把目标字典转换成 索引：字的形式\n",
    "tgt_vocab_size = len(tgt_vocab)                                                     # 目标字典尺寸\n",
    "\n",
    "src_len = len(sentences[0][0].split(\" \"))                                           # Encoder输入的最大长度 5\n",
    "tgt_len = len(sentences[0][1].split(\" \"))                                           # Decoder输入输出最大长度 5\n",
    "src_len,tgt_len\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 把sentences 转换成字典索引\n",
    "def make_data(sentences):\n",
    "    enc_inputs, dec_inputs, dec_outputs = [], [], []\n",
    "    for i in range(len(sentences)):\n",
    "      enc_input = [[src_vocab[n] for n in sentences[i][0].split()]] \n",
    "      dec_input = [[tgt_vocab[n] for n in sentences[i][1].split()]] \n",
    "      dec_output = [[tgt_vocab[n] for n in sentences[i][2].split()]] \n",
    "      enc_inputs.extend(enc_input)\n",
    "      dec_inputs.extend(dec_input)\n",
    "      dec_outputs.extend(dec_output)\n",
    "    print(enc_inputs,\"\\n   \",dec_inputs,\"\\n   \",dec_outputs)\n",
    "    return torch.LongTensor(enc_inputs), torch.LongTensor(dec_inputs), torch.LongTensor(dec_outputs)\n",
    "enc_inputs, dec_inputs, dec_outputs = make_data(sentences)\n",
    "print(enc_inputs)\n",
    "print(dec_inputs)\n",
    "print(dec_outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#自定义数据集函数\n",
    "class MyDataSet(Data.Dataset):\n",
    "  def __init__(self, enc_inputs, dec_inputs, dec_outputs):\n",
    "    super(MyDataSet, self).__init__()\n",
    "    self.enc_inputs = enc_inputs\n",
    "    self.dec_inputs = dec_inputs\n",
    "    self.dec_outputs = dec_outputs\n",
    "  \n",
    "  def __len__(self):\n",
    "    return self.enc_inputs.shape[0]\n",
    "  \n",
    "  def __getitem__(self, idx):\n",
    "    return self.enc_inputs[idx], self.dec_inputs[idx], self.dec_outputs[idx]\n",
    "loader = Data.DataLoader(MyDataSet(enc_inputs, dec_inputs, dec_outputs), 2, True) \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_model = 512   # 字 Embedding 的维度\n",
    "d_ff = 2048     # 前向传播隐藏层维度\n",
    "d_k = d_v = 64  # K(=Q), V的维度 \n",
    "n_layers = 6    # 有多少个encoder和decoder\n",
    "n_heads = 8     # Multi-Head Attention设置为8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self,d_model,dropout=0.1,max_len=5000):\n",
    "        super(PositionalEncoding,self).__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "        pos_table = np.array([\n",
    "        [pos / np.power(10000, 2 * i / d_model) for i in range(d_model)]\n",
    "        if pos != 0 else np.zeros(d_model) for pos in range(max_len)])\n",
    "        pos_table[1:, 0::2] = np.sin(pos_table[1:, 0::2])                  # 字嵌入维度为偶数时\n",
    "        pos_table[1:, 1::2] = np.cos(pos_table[1:, 1::2])                  # 字嵌入维度为奇数时\n",
    "        self.pos_table = torch.FloatTensor(pos_table)               # enc_inputs: [seq_len, d_model]\n",
    "    def forward(self,enc_inputs):                                         # enc_inputs: [batch_size, seq_len, d_model]\n",
    "        enc_inputs += self.pos_table[:enc_inputs.size(1),:]\n",
    "        return self.dropout(enc_inputs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_attn_pad_mask(seq_q,seq_k):\n",
    "    batch_size, len_q = seq_q.size()# seq_q 用于升维，为了做attention，mask score矩阵用的\n",
    "    batch_size, len_k = seq_k.size()\n",
    "    pad_attn_mask = seq_k.data.eq(0).unsqueeze(1) # 判断 输入那些含有P(=0),用1标记 ,[batch_size, 1, len_k]\n",
    "    return pad_attn_mask.expand(batch_size,len_q,len_k) # 扩展成多维度   [batch_size, len_q, len_k]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ScaledDotProductAttention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(ScaledDotProductAttention, self).__init__()\n",
    "    def forward(self, Q, K, V, attn_mask):                             # Q: [batch_size, n_heads, len_q, d_k]\n",
    "                                                                       # K: [batch_size, n_heads, len_k, d_k]\n",
    "                                                                       # V: [batch_size, n_heads, len_v(=len_k), d_v]\n",
    "                                                                       # attn_mask: [batch_size, n_heads, seq_len, seq_len]\n",
    "        scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k)   # scores : [batch_size, n_heads, len_q, len_k]\n",
    "        scores.masked_fill_(attn_mask, -1e9)                           # 如果是停用词P就等于 0 \n",
    "        attn = nn.Softmax(dim=-1)(scores)\n",
    "        context = torch.matmul(attn, V)                                # [batch_size, n_heads, len_q, d_v]\n",
    "        return context, attn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)\n",
    "        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)\n",
    "        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)\n",
    "        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)    \n",
    "    def forward(self, input_Q, input_K, input_V, attn_mask):    # input_Q: [batch_size, len_q, d_model]\n",
    "                                                                # input_K: [batch_size, len_k, d_model]\n",
    "                                                                # input_V: [batch_size, len_v(=len_k), d_model]\n",
    "                                                                # attn_mask: [batch_size, seq_len, seq_len]\n",
    "        residual, batch_size = input_Q, input_Q.size(0)\n",
    "        Q = self.W_Q(input_Q).view(batch_size, -1, n_heads, d_k).transpose(1,2)  # Q: [batch_size, n_heads, len_q, d_k]\n",
    "        K = self.W_K(input_K).view(batch_size, -1, n_heads, d_k).transpose(1,2)  # K: [batch_size, n_heads, len_k, d_k]\n",
    "        V = self.W_V(input_V).view(batch_size, -1, n_heads, d_v).transpose(1,2)  # V: [batch_size, n_heads, len_v(=len_k), d_v]\n",
    "        attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1)              # attn_mask : [batch_size, n_heads, seq_len, seq_len]\n",
    "        context, attn = ScaledDotProductAttention()(Q, K, V, attn_mask)          # context: [batch_size, n_heads, len_q, d_v]\n",
    "                                                                                 # attn: [batch_size, n_heads, len_q, len_k]\n",
    "        context = context.transpose(1, 2).reshape(batch_size, -1, n_heads * d_v) # context: [batch_size, len_q, n_heads * d_v]\n",
    "        output = self.fc(context)                                                # [batch_size, len_q, d_model]\n",
    "        return nn.LayerNorm(d_model)(output + residual), attn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PoswiseFeedForwardNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(PoswiseFeedForwardNet, self).__init__()\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(d_model, d_ff, bias=False),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(d_ff, d_model, bias=False))\n",
    "        \n",
    "    def forward(self, inputs):                             # inputs: [batch_size, seq_len, d_model]\n",
    "        residual = inputs\n",
    "        output = self.fc(inputs)\n",
    "        return nn.LayerNorm(d_model)(output + residual)   # [batch_size, seq_len, d_model]  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.enc_self_attn = MultiHeadAttention()                                     # 多头注意力机制\n",
    "        self.pos_ffn = PoswiseFeedForwardNet()                                        # 前馈神经网络\n",
    "    def forward(self, enc_inputs, enc_self_attn_mask):                                # enc_inputs: [batch_size, src_len, d_model]\n",
    "        #输入3个enc_inputs分别与W_q、W_k、W_v相乘得到Q、K、V                          # enc_self_attn_mask: [batch_size, src_len, src_len]\n",
    "        enc_outputs, attn = self.enc_self_attn(enc_inputs, enc_inputs, enc_inputs,    # enc_outputs: [batch_size, src_len, d_model], \n",
    "                                               enc_self_attn_mask)                    # attn: [batch_size, n_heads, src_len, src_len]                                                                   \n",
    "        enc_outputs = self.pos_ffn(enc_outputs)                                       # enc_outputs: [batch_size, src_len, d_model]\n",
    "        return enc_outputs, attn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.src_emb = nn.Embedding(src_vocab_size, d_model)\n",
    "        self.pos_emb = PositionalEncoding(d_model)\n",
    "        self.layers = nn.ModuleList([EncoderLayer() for _ in range(n_layers)])\n",
    "    def forward(self, enc_inputs):\n",
    "        '''\n",
    "        enc_inputs: [batch_size, src_len]\n",
    "        '''\n",
    "        enc_outputs = self.src_emb(enc_inputs) # [batch_size, src_len, d_model]\n",
    "        enc_outputs = self.pos_emb(enc_outputs.transpose(0, 1)).transpose(0, 1) # [batch_size, src_len, d_model]\n",
    "        enc_self_attn_mask = get_attn_pad_mask(enc_inputs, enc_inputs) # [batch_size, src_len, src_len]\n",
    "        enc_self_attns = []\n",
    "        for layer in self.layers:\n",
    "            # enc_outputs: [batch_size, src_len, d_model], enc_self_attn: [batch_size, n_heads, src_len, src_len]\n",
    "            enc_outputs, enc_self_attn = layer(enc_outputs, enc_self_attn_mask)\n",
    "            enc_self_attns.append(enc_self_attn)\n",
    "        return enc_outputs, enc_self_attns\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.dec_self_attn = MultiHeadAttention()\n",
    "        self.dec_enc_attn = MultiHeadAttention()\n",
    "        self.pos_ffn = PoswiseFeedForwardNet()\n",
    "    def forward(self, dec_inputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask): # dec_inputs: [batch_size, tgt_len, d_model]\n",
    "                                                                                       # enc_outputs: [batch_size, src_len, d_model]\n",
    "                                                                                       # dec_self_attn_mask: [batch_size, tgt_len, tgt_len]\n",
    "                                                                                       # dec_enc_attn_mask: [batch_size, tgt_len, src_len]\n",
    "        dec_outputs, dec_self_attn = self.dec_self_attn(dec_inputs, dec_inputs, \n",
    "                                                 dec_inputs, dec_self_attn_mask)   # dec_outputs: [batch_size, tgt_len, d_model]\n",
    "                                                                                   # dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len]\n",
    "        dec_outputs, dec_enc_attn = self.dec_enc_attn(dec_outputs, enc_outputs, \n",
    "                                                enc_outputs, dec_enc_attn_mask)    # dec_outputs: [batch_size, tgt_len, d_model]\n",
    "                                                                                   # dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]\n",
    "        dec_outputs = self.pos_ffn(dec_outputs)                                    # dec_outputs: [batch_size, tgt_len, d_model]\n",
    "        return dec_outputs, dec_self_attn, dec_enc_attn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_attn_subsequence_mask(seq):                               # seq: [batch_size, tgt_len]\n",
    "    attn_shape = [seq.size(0), seq.size(1), seq.size(1)]          # 生成上三角矩阵,[batch_size, tgt_len, tgt_len]\n",
    "    subsequence_mask = np.triu(np.ones(attn_shape), k=1)\n",
    "    subsequence_mask = torch.from_numpy(subsequence_mask).byte()  #  [batch_size, tgt_len, tgt_len]\n",
    "    return subsequence_mask\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.tgt_emb = nn.Embedding(tgt_vocab_size, d_model)\n",
    "        self.pos_emb = PositionalEncoding(d_model)\n",
    "        self.layers = nn.ModuleList([DecoderLayer() for _ in range(n_layers)])\n",
    "    def forward(self, dec_inputs, enc_inputs, enc_outputs):\n",
    "        '''\n",
    "        dec_inputs: [batch_size, tgt_len]\n",
    "        enc_intpus: [batch_size, src_len]\n",
    "        enc_outputs: [batch_size, src_len, d_model]\n",
    "        '''\n",
    "        dec_outputs = self.tgt_emb(dec_inputs) # [batch_size, tgt_len, d_model]\n",
    "        dec_outputs = self.pos_emb(dec_outputs.transpose(0, 1)).transpose(0, 1) # [batch_size, tgt_len, d_model]\n",
    "        # Decoder输入序列的pad mask矩阵（这个例子中decoder是没有加pad的，实际应用中都是有pad填充的）\n",
    "        dec_self_attn_pad_mask = get_attn_pad_mask(dec_inputs, dec_inputs) # [batch_size, tgt_len, tgt_len]\n",
    "        # Masked Self_Attention：当前时刻是看不到未来的信息的\n",
    "        dec_self_attn_subsequence_mask = get_attn_subsequence_mask(dec_inputs) # [batch_size, tgt_len, tgt_len]\n",
    "        # Decoder中把两种mask矩阵相加（既屏蔽了pad的信息，也屏蔽了未来时刻的信息）\n",
    "        dec_self_attn_mask = torch.gt((dec_self_attn_pad_mask + dec_self_attn_subsequence_mask), 0) # [batch_size, tgt_len, tgt_len]\n",
    "        # 这个mask主要用于encoder-decoder attention层\n",
    "        # get_attn_pad_mask主要是enc_inputs的pad mask矩阵(因为enc是处理K,V的，求Attention时是用v1,v2,..vm去加权的，\n",
    "        # 要把pad对应的v_i的相关系数设为0，这样注意力就不会关注pad向量)\n",
    "        #                       dec_inputs只是提供expand的size的\n",
    "        dec_enc_attn_mask = get_attn_pad_mask(dec_inputs, enc_inputs) # [batc_size, tgt_len, src_len]\n",
    "        dec_self_attns, dec_enc_attns = [], []\n",
    "        for layer in self.layers:\n",
    "            # dec_outputs: [batch_size, tgt_len, d_model], dec_self_attn: [batch_size, n_heads, tgt_len, tgt_len], dec_enc_attn: [batch_size, h_heads, tgt_len, src_len]\n",
    "            dec_outputs, dec_self_attn, dec_enc_attn = layer(dec_outputs, enc_outputs, dec_self_attn_mask, dec_enc_attn_mask)\n",
    "            dec_self_attns.append(dec_self_attn)\n",
    "            dec_enc_attns.append(dec_enc_attn)\n",
    "        return dec_outputs, dec_self_attns, dec_enc_attns\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Transformer, self).__init__()\n",
    "        self.Encoder = Encoder()\n",
    "        self.Decoder = Decoder()\n",
    "        self.projection = nn.Linear(d_model, tgt_vocab_size, bias=False)\n",
    "    def forward(self, enc_inputs, dec_inputs):                         # enc_inputs: [batch_size, src_len]  \n",
    "                                                                       # dec_inputs: [batch_size, tgt_len]\n",
    "        enc_outputs, enc_self_attns = self.Encoder(enc_inputs)         # enc_outputs: [batch_size, src_len, d_model], \n",
    "                                                                       # enc_self_attns: [n_layers, batch_size, n_heads, src_len, src_len]\n",
    "        dec_outputs, dec_self_attns, dec_enc_attns = self.Decoder(\n",
    "            dec_inputs, enc_inputs, enc_outputs)                       # dec_outpus    : [batch_size, tgt_len, d_model], \n",
    "                                                                       # dec_self_attns: [n_layers, batch_size, n_heads, tgt_len, tgt_len], \n",
    "                                                                       # dec_enc_attn  : [n_layers, batch_size, tgt_len, src_len]\n",
    "        dec_logits = self.projection(dec_outputs)                      # dec_logits: [batch_size, tgt_len, tgt_vocab_size]\n",
    "        return dec_logits.view(-1, dec_logits.size(-1)), enc_self_attns, dec_self_attns, dec_enc_attns\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Transformer()\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=0)     #忽略 占位符 索引为0.\n",
    "optimizer = optim.SGD(model.parameters(), lr=1e-3, momentum=0.99)\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"模型参数总数: {total_params}\",\"device:\",device)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0001 loss = 0.001256\n",
      "Epoch: 0001 loss = 0.104630\n",
      "Epoch: 0001 loss = 0.002480\n",
      "Epoch: 0002 loss = 0.256928\n",
      "Epoch: 0002 loss = 0.001250\n",
      "Epoch: 0002 loss = 0.001297\n",
      "Epoch: 0003 loss = 0.001124\n",
      "Epoch: 0003 loss = 0.015961\n",
      "Epoch: 0003 loss = 0.006783\n",
      "Epoch: 0004 loss = 0.001159\n",
      "Epoch: 0004 loss = 0.036861\n",
      "Epoch: 0004 loss = 0.009819\n",
      "Epoch: 0005 loss = 0.007283\n",
      "Epoch: 0005 loss = 0.145971\n",
      "Epoch: 0005 loss = 0.001049\n",
      "Epoch: 0006 loss = 0.003890\n",
      "Epoch: 0006 loss = 0.024921\n",
      "Epoch: 0006 loss = 0.001239\n",
      "Epoch: 0007 loss = 0.000386\n",
      "Epoch: 0007 loss = 0.006563\n",
      "Epoch: 0007 loss = 0.000962\n",
      "Epoch: 0008 loss = 0.000759\n",
      "Epoch: 0008 loss = 0.001376\n",
      "Epoch: 0008 loss = 0.029969\n",
      "Epoch: 0009 loss = 0.000294\n",
      "Epoch: 0009 loss = 0.102409\n",
      "Epoch: 0009 loss = 0.003785\n",
      "Epoch: 0010 loss = 0.218084\n",
      "Epoch: 0010 loss = 0.009162\n",
      "Epoch: 0010 loss = 0.001629\n",
      "Epoch: 0011 loss = 0.001088\n",
      "Epoch: 0011 loss = 0.000537\n",
      "Epoch: 0011 loss = 0.124538\n",
      "Epoch: 0012 loss = 0.821644\n",
      "Epoch: 0012 loss = 0.000521\n",
      "Epoch: 0012 loss = 0.001338\n",
      "Epoch: 0013 loss = 0.000598\n",
      "Epoch: 0013 loss = 0.000267\n",
      "Epoch: 0013 loss = 0.349878\n",
      "Epoch: 0014 loss = 0.000527\n",
      "Epoch: 0014 loss = 0.008569\n",
      "Epoch: 0014 loss = 0.000521\n",
      "Epoch: 0015 loss = 0.000440\n",
      "Epoch: 0015 loss = 0.000293\n",
      "Epoch: 0015 loss = 0.725728\n",
      "Epoch: 0016 loss = 0.000458\n",
      "Epoch: 0016 loss = 0.580913\n",
      "Epoch: 0016 loss = 0.000397\n",
      "Epoch: 0017 loss = 0.001772\n",
      "Epoch: 0017 loss = 0.000844\n",
      "Epoch: 0017 loss = 0.064775\n",
      "Epoch: 0018 loss = 0.000174\n",
      "Epoch: 0018 loss = 0.002537\n",
      "Epoch: 0018 loss = 0.002440\n",
      "Epoch: 0019 loss = 0.009115\n",
      "Epoch: 0019 loss = 0.000878\n",
      "Epoch: 0019 loss = 0.044293\n",
      "Epoch: 0020 loss = 0.110389\n",
      "Epoch: 0020 loss = 0.000446\n",
      "Epoch: 0020 loss = 0.000512\n"
     ]
    }
   ],
   "source": [
    "# 将模型设置为训练模式\n",
    "model.train()\n",
    "for epoch in range(20):\n",
    "    for enc_inputs, dec_inputs, dec_outputs in loader:\n",
    "        enc_inputs, dec_inputs, dec_outputs = enc_inputs, dec_inputs, dec_outputs\n",
    "        outputs, enc_self_attns, dec_self_attns, dec_enc_attns = model(enc_inputs, dec_inputs)        \n",
    "        loss = criterion(outputs,dec_outputs.view(-1))\n",
    "        #print(outputs,'\\n\\n\\n\\n\\n',dec_outputs.view(-1))\n",
    "        print('Epoch:', '%04d' % (epoch+1), 'loss =', '{:.6f}'.format(loss))\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 保存模型状态字典、优化器状态字典和训练轮数\n",
    "torch.save({\n",
    "    'model_state_dict': model.state_dict(),\n",
    "    'optimizer_state_dict': optimizer.state_dict(),\n",
    "    'epoch': epoch\n",
    "}, './model.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 加载模型状态字典、优化器状态字典和训练轮数\n",
    "checkpoint = torch.load('./model.pth')\n",
    "model.load_state_dict(checkpoint['model_state_dict'])\n",
    "optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "epoch = checkpoint['epoch']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['我', '是', '男', '生', 'P'] -> ['I', 'am', 'a', 'boy', 'E']\n"
     ]
    }
   ],
   "source": [
    "# 将模型设置为评估模式\n",
    "model.eval()\n",
    "def test(model, enc_input, start_symbol):\n",
    "    enc_outputs, enc_self_attns = model.Encoder(enc_input)\n",
    "    dec_input = torch.zeros(1,tgt_len).type_as(enc_input.data)\n",
    "    next_symbol = start_symbol\n",
    "    for i in range(0,tgt_len):\n",
    "        dec_input[0][i] = next_symbol\n",
    "        dec_outputs, _, _ = model.Decoder(dec_input,enc_input,enc_outputs)\n",
    "        projected = model.projection(dec_outputs)\n",
    "        prob = projected.squeeze(0).max(dim=-1, keepdim=False)[1]\n",
    "        next_word = prob.data[i]\n",
    "        next_symbol = next_word.item()\n",
    "    return dec_input\n",
    "enc_inputs, _, _ = next(iter(loader))\n",
    "predict_dec_input = test(model, enc_inputs[1].view(1, -1), start_symbol=tgt_vocab[\"S\"])\n",
    "predict, _, _, _ = model(enc_inputs[1].view(1, -1), predict_dec_input)\n",
    "predict = predict.data.max(1, keepdim=True)[1]\n",
    "print([src_idx2word[int(i)] for i in enc_inputs[1]], '->', \n",
    "[idx2word[n.item()] for n in predict.squeeze()])\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
